from __future__ import print_function, division
import os
from PIL import Image
import numpy as np
from torch.utils.data import Dataset
from mypath import Path
from torchvision import transforms
from dataloaders import custom_transforms as tr



class VOCSegmentation(Dataset):
    """
    PascalVoc dataset
    """
    NUM_CLASSES = 21

    def __init__(self,
                 args,                                  # 配置参数
                 base_dir=Path.db_root_dir('pascal'),   # 数据集路径
                 split='train',                         # 数据模式
                 ):
        """
        :param base_dir: path to VOC dataset directory
        :param split: train/val
        :param transform: transform to apply
        """
        super().__init__()
        self._base_dir = base_dir
        self._image_dir = os.path.join(self._base_dir, 'JPEGImages')
        self._cat_dir = os.path.join(self._base_dir, 'SegmentationClass')

        if isinstance(split, str):
            self.split = [split]
        else:
            split.sort()
            self.split = split

        self.args = args

        _splits_dir = os.path.join(self._base_dir, 'ImageSets', 'Segmentation')

        self.im_ids = []
        self.images = []
        self.categories = []

        for splt in self.split:
            with open(os.path.join(os.path.join(_splits_dir, splt + '.txt')), "r") as f:
                lines = f.read().splitlines()

            for ii, line in enumerate(lines):
                _image = os.path.join(self._image_dir, line + ".jpg")
                _cat = os.path.join(self._cat_dir, line + ".png")
                assert os.path.isfile(_image)
                assert os.path.isfile(_cat)
                self.im_ids.append(line)
                self.images.append(_image)
                self.categories.append(_cat)

        assert (len(self.images) == len(self.categories))

        # Display stats
        print('Number of images in {}: {:d}'.format(split, len(self.images)))

    def __len__(self):
        return len(self.images)

    def __getitem__(self, index):
        _img, _target = self._make_img_gt_point_pair(index)
        sample = {'image': _img, 'label': _target}

        for split in self.split:
            if split == "train":
                return self.transform_tr(sample)
            elif split == 'val':
                return self.transform_val(sample)

    def _make_img_gt_point_pair(self, index):
        _img = Image.open(self.images[index]).convert('RGB')
        _target = Image.open(self.categories[index])

        return _img, _target

    def transform_tr(self, sample):
        composed_transforms = transforms.Compose([
            tr.RandomHorizontalFlip(),
            tr.RandomScaleCrop(base_size=self.args.base_size, crop_size=self.args.crop_size),
            tr.RandomGaussianBlur(),
            tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
            tr.ToTensor()])

        return composed_transforms(sample)

    def transform_val(self, sample):

        composed_transforms = transforms.Compose([
            tr.FixScaleCrop(crop_size=self.args.crop_size),
            tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
            tr.ToTensor()])

        return composed_transforms(sample)

    def __str__(self):
        return 'VOC2012(split=' + str(self.split) + ')'


def my_transform_tr(sample):
    composed_transforms = transforms.Compose([
        tr.RandomHorizontalFlip(),  # 左右
        tr.RandomHorizontalFlip2(),  # 上下
        tr.RandomRotate(10),  # 旋转
        # tr.RandomScaleCrop(base_size=513, crop_size=256),
        tr.RandomCrop_50_100(crop_size=256),
        tr.RandomGaussianBlur(),
        # tr.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
        tr.ToTensor()])

    return composed_transforms(sample)


if __name__ == '__main__':
    from torch.utils.data import DataLoader
    import matplotlib.pyplot as plt
    import argparse

    parser = argparse.ArgumentParser()
    args = parser.parse_args()
    args.base_size = 256
    args.crop_size = 256

    path1 = r'E:\Dataset\Pascal_Voc\VOCdevkit\VOC2012\JPEGImages\2007_001408.jpg'
    path2 = r'E:\Dataset\Pascal_Voc\VOCdevkit\VOC2012\SegmentationClass\2007_001408.png'

    plt.figure()
    plt.title('display')
    # 500*333
    _img = Image.open(path1).convert('RGB')
    _target = Image.open(path2)
    print(_img.size)
    sample = {'image': _img, 'label': _target}
    plt.subplot(221)
    plt.imshow(_img)
    plt.subplot(222)
    plt.imshow(_target)
    plt.show(block=True)
    exit()

    sample = my_transform_tr(sample)

    img = sample['image'].numpy()
    gt = sample['label'].numpy()
    print(type(img))
    print(type(gt))
    print(img.size)
    print(gt.size)
    print(img.shape)
    print(gt.shape)
    tmp = np.array(gt).astype(np.uint8)
    # segmap = decode_segmap(tmp, dataset='pascal')

    img_tmp = np.transpose(img, axes=[1, 2, 0])
    # img_tmp *= (0.229, 0.224, 0.225)
    # img_tmp += (0.485, 0.456, 0.406)
    # img_tmp *= 255.0
    print(type(img))
    print(img.max())
    img_tmp = img_tmp.astype(np.uint8)
    # print(_img.shape)
    # segmap = decode_segmap(_target, dataset='pascal')

    plt.subplot(223)
    plt.imshow(img_tmp)
    plt.subplot(224)
    plt.imshow(segmap)
    plt.show(block=True)

    exit()

    voc_train = VOCSegmentation(args, split='train')
    dataloader = DataLoader(voc_train, batch_size=5, shuffle=False, num_workers=2)

    for ii, sample in enumerate(dataloader):
        for jj in range(sample["image"].size()[0]):
            img = sample['image'].numpy()
            gt = sample['label'].numpy()
            tmp = np.array(gt[jj]).astype(np.uint8)
            # segmap = decode_segmap(tmp, dataset='pascal')
            img_tmp = np.transpose(img[jj], axes=[1, 2, 0])
            img_tmp *= (0.229, 0.224, 0.225)
            img_tmp += (0.485, 0.456, 0.406)
            img_tmp *= 255.0
            img_tmp = img_tmp.astype(np.uint8)
            plt.figure()
            plt.title('display')
            plt.subplot(211)
            plt.imshow(img_tmp)
            plt.subplot(212)
            # plt.imshow(segmap)

        if ii == 0:
            break

    plt.show(block=True)


